import pymongo
from common.dir_config import ConfigHandler
from util.readFileUtils.yamlControl import GetYamlData
import socket


def on_timeout_callback(data, retries=3):
    print("MongoDB server timeout, retrying...")
    for i in range(retries):
        try:
            # 重新连接
            client = pymongo.MongoClient(data, serverSelectionTimeoutMS=5000)
            client.set_socket_timeout_ms(5000)
            return client
        except socket.timeout as e:
            print("Retry connecting to MongoDB server ({}/{})".format(i + 1, retries))
    print("MongoDB connection failed!")
    return None


def con_mangodb(collet_name, table_name, evn, address, appoint_fields=None,**query,):
    ''' 连接mongodb地址host,端口port（如果不给它传递参数，默认是27017）
     evn:1是测试环境，2是生产环境
     address:连接地址'''

    if evn == 1:
        data = GetYamlData(ConfigHandler.test_yaml_dir).get_yaml_data()
    elif evn == 2:
        data = GetYamlData(ConfigHandler.pro_yaml_dir).get_yaml_data()
    try:
        data = data['amazon_mongodb'][address]
        # print(data)
        con = pymongo.MongoClient(data, serverSelectionTimeoutMS=5000)
        # con.set_socket_timeout_ms(5000)
        # con.on_timeout = on_timeout_callback(data)
        mycol1 = con[collet_name][table_name]
        # # 查询条件26743
        myquery = query
        # # 查询想要的数据
        if appoint_fields:
            result = mycol1.find(myquery).limit(10000).skip(0)
        else:
            result = mycol1.find(myquery, appoint_fields).limit(10000).skip(0)
        results = []
        for i in result:
            results.append(i)
    except Exception as e:
        # 捕获超时异常
        print("MongoDB connection timeout: ", e)
    return results


if __name__ == '__main__':
    '''print(con_mangodb("Amazon_AdjustPrice", "AdjustPriceQueue", 2,
                      'url1', OrderSourceID=16405))'''
